#include "kernel_operator.h"

using namespace AscendC;

#define ll long long

constexpr int32_t BUFFER_NUM = 1;

template<typename T> class GcdBrute {
private:
    TPipe* pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<T> x1Gm;
    GlobalTensor<T> x2Gm;
    GlobalTensor<T> yGm;
    int totalLength1;
    int totalLength2;
    int totalLength;
    int dimNum1;
    int dimNum2;
    int dimNum;
    int* dim1;
    int* dim2;
    int* dim;

public:
    __aicore__ inline GcdBrute() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, int totalLength1, int totalLength2, int totalLength, int dimNum1, int dimNum2, int dimNum, int* dim1, int* dim2, int* dim)
    {
        pipe = pipeIn;
        this->totalLength1 = totalLength1;
        this->totalLength2 = totalLength2;
        this->totalLength = totalLength;
        this->dimNum1 = dimNum1;
        this->dimNum2 = dimNum2;
        this->dimNum = dimNum;
        this->dim1 = dim1;
        this->dim2 = dim2;
        this->dim = dim;

        x1Gm.SetGlobalBuffer((__gm__ T *)x1, this->totalLength1);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2, this->totalLength2);
        yGm.SetGlobalBuffer((__gm__ T *)y, this->totalLength);

        // pipe->InitBuffer(inQueueX1, BUFFER_NUM, this->totalLength1 * sizeof(T));
        // pipe->InitBuffer(inQueueX2, BUFFER_NUM, this->totalLength2 * sizeof(T));
        // pipe->InitBuffer(outQueueY, BUFFER_NUM, this->totalLength * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        // 假设所有输入具有相同的维度数量
        // 创建索引数组记录当前处理的多维坐标
        int indices[5] = {0};
        int count = 0;
        
        // 处理所有元素
        while (count < totalLength) {
            // 计算x1索引
            int x1_idx = 0;
            int stride1 = 1;
            for (int i = dimNum - 1; i >= 0; i--) {
                int idx = indices[i] % dim1[i]; // 处理广播
                x1_idx += idx * stride1;
                stride1 *= dim1[i];
            }
            
            // 计算x2索引
            int x2_idx = 0;
            int stride2 = 1;
            for (int i = dimNum - 1; i >= 0; i--) {
                int idx = indices[i] % dim2[i]; // 处理广播
                x2_idx += idx * stride2;
                stride2 *= dim2[i];
            }
            
            // 获取x1和x2值并计算GCD
            ll x1 = x1Gm(x1_idx);
            ll x2 = x2Gm(x2_idx);
            x1 = x1 < 0 ? -x1 : x1;
            x2 = x2 < 0 ? -x2 : x2;
            ll y = gcd(x1, x2);
            yGm(count) = y;

            // if (count < 10 || count + 10 > totalLength) {
            //     // 打印所有东西
            //     printf("count: %d, x1_idx: %d, x2_idx: %d, x1: %lld, x2: %lld, y: %lld\n", count, x1_idx, x2_idx, x1, x2, y);
            //     // indices
            //     printf("indices: ");
            //     for (int i = 0; i < dimNum; i++) {
            //         printf("%d ", indices[i]);
            //     }
            //     printf("\n");
            // }

            // 更新indices数组，类似多维数组的索引递增
            count++;
            int pos = dimNum - 1;
            while (pos >= 0) {
                indices[pos]++;
                if (indices[pos] < dim[pos]) {
                    break;
                }
                indices[pos] = 0;
                pos--;
            }
        }
    }

private:
    __aicore__ inline ll gcd(ll a, ll b)
    {
        while (b != 0) {
            ll t = b;
            b = a % b;
            a = t;
        }
        return a;
    }
};

template<typename T> class Gcd {
private:
    TPipe* pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<T> x1Gm;
    GlobalTensor<T> x2Gm;
    GlobalTensor<T> yGm;
    int blockIdx;
    int offset;
    int offset1;
    int offset2;
    int indices[5];
    int totalLength1;
    int totalLength2;
    int totalLength;
    int calcLength;
    int dimNum1;
    int dimNum2;
    int dimNum;
    int* dim1;
    int* dim2;
    int* dim;

public:
    __aicore__ inline Gcd() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, int totalLength1, int totalLength2, int totalLength, int dimNum1, int dimNum2, int dimNum, int* dim1, int* dim2, int* dim)
    {
        pipe = pipeIn;
        this->totalLength1 = totalLength1;
        this->totalLength2 = totalLength2;
        this->totalLength = totalLength;
        this->dimNum1 = dimNum1;
        this->dimNum2 = dimNum2;
        this->dimNum = dimNum;
        this->dim1 = dim1;
        this->dim2 = dim2;
        this->dim = dim;

        blockIdx = GetBlockIdx();
        calcLength = dim[dimNum - 1];
        offset1 = blockIdx * calcLength;
        // 计算indices（每一维的索引）
        indices[dimNum - 1] = 0;
        int tmp = blockIdx;
        for (int i = dimNum - 2; i >= 0; i--) {
            indices[i] = tmp % dim[i];
            tmp /= dim[i];
        }
        // printf("blockIdx: %d, indices: ", blockIdx);
        // for (int i = 0; i < dimNum; i++) {
        //     printf("%d ", indices[i]);
        // }
        // printf("\n");

        // 计算x2索引
        offset2 = 0;
        int stride2 = 1;
        for (int i = dimNum - 1; i >= 0; i--) {
            int idx = indices[i] % dim2[i]; // 处理广播
            offset2 += idx * stride2;
            stride2 *= dim2[i];
        }
        // printf("offset1: %d, offset2: %d\n", offset1, offset2);
        

        x1Gm.SetGlobalBuffer((__gm__ T *)x1 + offset1, calcLength);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2 + offset2, dim2[dimNum - 1]);
        yGm.SetGlobalBuffer((__gm__ T *)y + offset1, calcLength);

        pipe->InitBuffer(inQueueX1, BUFFER_NUM, calcLength * sizeof(T));
        pipe->InitBuffer(inQueueX2, BUFFER_NUM, calcLength * sizeof(T));
        pipe->InitBuffer(outQueueY, BUFFER_NUM, calcLength * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        // if (dimNum < 3) {
        //     pipe->InitBuffer(inQueueX1, BUFFER_NUM, 100000000);
        // } else if (dim1[dimNum - 3] != dim2[dimNum - 3]) {
        //     return;
        // }
        // for (int i = 0; i < dim[dimNum - 1]; i++) {
        //     // 获取x1和x2值并计算GCD
        //     int x1 = x1Gm(i);
        //     int x2 = dim2[dimNum - 1] == 1 ? x2Gm(0) : x2Gm(i);
        //     x1 = x1 < 0 ? -x1 : x1;
        //     x2 = x2 < 0 ? -x2 : x2;
        //     int y = gcd(x1, x2);
        //     yGm(i) = y;
        // }
        DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        // VECIN和VECOUT是32B单位，其它都是1B单位
        DataCopyExtParams copyParams = {
            (uint16_t)1, 
            (uint32_t)(calcLength * sizeof(T)), 
            0,
            0,
            0
        };
        DataCopyExtParams copyParams2 = {
            (uint16_t)1, 
            (uint32_t)(1 * sizeof(T)), 
            0,
            0,
            0
        };
            // 处理x2为多维的情况
            LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
            LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
            // DataCopy(x1Local, x1Gm, calcLength);
            // DataCopy(x2Local, x2Gm, calcLength);
            DataCopyPad(x1Local, x1Gm, copyParams, padParams);
            DataCopyPad(x2Local, x2Gm, copyParams, padParams);
            inQueueX1.EnQue(x1Local);
            inQueueX2.EnQue(x2Local);
            x1Local = inQueueX1.DeQue<T>();
            x2Local = inQueueX2.DeQue<T>();
            LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();
            for (int i = 0; i < calcLength; i++) {
                T x1 = x1Local(i);
                T x2 = x2Local(i);
                x1 = x1 < 0 ? -x1 : x1;
                x2 = x2 < 0 ? -x2 : x2;
                T y = gcd(x1, x2);
                yLocal(i) = y;
            }
            outQueueY.EnQue(yLocal);
            inQueueX1.FreeTensor(x1Local);
            inQueueX2.FreeTensor(x2Local);
            yLocal = outQueueY.DeQue<T>();
            // DataCopy(yGm, yLocal, calcLength);
            DataCopyPad(yGm, yLocal, copyParams);
            outQueueY.FreeTensor(yLocal);

        // if (count < 10 || count + 10 > totalLength) {
        //     // 打印所有东西
        //     printf("count: %d, x1_idx: %d, x2_idx: %d, x1: %lld, x2: %lld, y: %lld\n", count, x1_idx, x2_idx, x1, x2, y);
        //     // indices
        //     printf("indices: ");
        //     for (int i = 0; i < dimNum; i++) {
        //         printf("%d ", indices[i]);
        //     }
        //     printf("\n");
        // }
    }

private:
    __aicore__ inline T gcd(T a, T b)
    {
        while (b != 0) {
            T t = b;
            b = a % b;
            a = t;
        }
        return a;
    }
};

extern "C" __global__ __aicore__ void gcd(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    if (TILING_KEY_IS(0)) {
        // 暴力
        GcdBrute<DTYPE_X1> op;
        op.Init(x1, x2, y, workspace, &pipe, tiling_data.totalLength1, tiling_data.totalLength2, tiling_data.totalLength, tiling_data.dimNum1, tiling_data.dimNum2, tiling_data.dimNum, tiling_data.dim1, tiling_data.dim2, tiling_data.dim);
        op.Process();
    } else if (TILING_KEY_IS(1)) {
        // 每个核处理最后一维
        Gcd<DTYPE_X1> op;
        op.Init(x1, x2, y, workspace, &pipe, tiling_data.totalLength1, tiling_data.totalLength2, tiling_data.totalLength, tiling_data.dimNum1, tiling_data.dimNum2, tiling_data.dimNum, tiling_data.dim1, tiling_data.dim2, tiling_data.dim);
        op.Process();
    }
}